#!/usr/bin/env python3
# 导入所需的库和模块
# 从增强版多代理工作流导入所需类
from src.research_core.enhanced_multi_agent_workflow import create_enhanced_multi_agent_workflow, EnhancedMultiAgentState as MultiAgentState
from src.config import research_settings
from langchain_core.runnables import RunnableConfig
import json
import os
from datetime import datetime
import time
import asyncio

class InteractiveMultiAgent:
    """交互式多代理系统"""
    
    def __init__(self):
        self.sessions = {}  # 存储多个会话
        self.current_session_id = None
        self.session_history = []
        self.create_new_session()
        
    def create_new_session(self):
        """创建新会话"""
        session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.sessions[session_id] = {
            "history": [],
            "created_at": datetime.now().isoformat(),
            "name": f"会话_{len(self.sessions) + 1}"
        }
        self.current_session_id = session_id
        self.session_history = self.sessions[session_id]["history"]
        return session_id
        
    def switch_session(self, session_id):
        """切换到指定会话"""
        if session_id in self.sessions:
            self.current_session_id = session_id
            self.session_history = self.sessions[session_id]["history"]
            return True
        return False
        
    def list_sessions(self):
        """列出所有会话"""
        return self.sessions
        
    def rename_session(self, session_id, new_name):
        """重命名会话"""
        if session_id in self.sessions:
            self.sessions[session_id]["name"] = new_name
            return True
        return False
        
    def delete_session(self, session_id):
        """删除会话"""
        if session_id in self.sessions and len(self.sessions) > 1:
            del self.sessions[session_id]
            if self.current_session_id == session_id:
                # 切换到第一个可用会话
                self.current_session_id = next(iter(self.sessions.keys()))
                self.session_history = self.sessions[self.current_session_id]["history"]
            return True
        return False
        
    async def run_workflow(self, question):
        """运行多代理工作流处理单个问题"""
        start_time = time.time()
        print(f"[DEBUG] 开始处理问题: {question}")
        
        # 获取多代理工作流
        print("[DEBUG] 正在获取工作流...")
        try:
            # 直接创建增强版多代理工作流
            app = create_enhanced_multi_agent_workflow()
        except Exception as e:
            print(f"[DEBUG] 工作流创建失败: {e}")
            return None
        print(f"[DEBUG] 工作流获取完成: {type(app)}")

        # 定义初始状态（使用多代理状态）
        initial_state: MultiAgentState = {
            "question": question,
            "search_query": "",
            "search_results": "",
            "research_complete": False,
            "final_answer": "",
            "images": [],                    # List[dict]
            "tables": [],                    # List[dict]
            "multimodal_content": {},        # dict
            "search_strategy": "",           # 添加搜索策略字段
            "analysis_results": {},          # 添加分析结果字段
            "agent_assignments": {},         # 添加代理分配字段
            "iteration_count": 0,            # 添加迭代计数器
            "previous_states": [],           # 添加状态历史记录
            "cache_key": None,               # 添加缓存键
            "reflection": None,              # 添加质量评估结果
            "visualizations": {},            # 添加数据可视化结果
        }

        print(f"\n🚀 处理问题: {question}")
        print("=" * 60)
        
        # 定义配置，包括线程ID和递归限制
        config: RunnableConfig = {
            "recursion_limit": 50,  # 增加递归限制以允许更多迭代，从10增加到50
            "configurable": {
                "thread_id": "interactive_multi_agent_thread"
            }
        }
        
        # 初始化final_state变量，以防出现异常
        final_state = initial_state.copy()
        
        try:
            # 添加超时机制，最多运行2分钟
            print("[DEBUG] 开始执行工作流...")
            final_state = await app.ainvoke(initial_state, config=config)
            print(f"[DEBUG] 工作流执行完成，最终状态: {list(final_state.keys())}")
            print(f"[DEBUG] 最终答案: {final_state.get('final_answer', '无答案')[:100] if final_state.get('final_answer') else '无答案'}")
        except TimeoutError:
            print("❌ 工作流执行超时（2分钟）")
            return None
        except Exception as e:
            print(f"❌ 工作流执行出错: {str(e)}")
            import traceback
            traceback.print_exc()
        
        end_time = time.time()
        processing_time = end_time - start_time

        # 记录会话历史
        session_entry = {
            "timestamp": datetime.now().isoformat(),
            "question": question,
            "answer": final_state.get('final_answer', ''),
            "processing_time": round(processing_time, 2),
            "search_strategy": final_state.get('search_strategy', ''),
            "research_complete": final_state.get('research_complete', False)
        }
        self.session_history.append(session_entry)

        # 输出结果
        print("\n🎯 最终答案")
        print("=" * 60)
        answer = final_state.get('final_answer', '未生成答案')
        # 修复输出文本编码问题
        if isinstance(answer, str):
            answer = answer.encode('utf-8', errors='ignore').decode('utf-8')
        print(f"答案内容: {answer}")
        if not answer or answer == '未生成答案':
            print("⚠️  未生成有效答案，显示完整最终状态:")
            for key, value in final_state.items():
                print(f"  {key}: {str(value)[:100]}{'...' if len(str(value)) > 100 else ''}")

        # 查看中间步骤（只显示关键步骤，避免重复执行工作流）
        print("\n🔍 多代理协作过程")
        print("=" * 60)
        if final_state.get('search_strategy'):
            strategy = final_state['search_strategy']
            if isinstance(strategy, str):
                strategy = strategy.encode('utf-8', errors='ignore').decode('utf-8')
            print(f"📋 搜索策略专家: {strategy[:100]}...")
        if final_state.get('search_results'):
            results = final_state['search_results']
            if isinstance(results, str):
                results = results.encode('utf-8', errors='ignore').decode('utf-8')
            print(f"🔎 搜索专家: {results[:100]}...")
        print("=" * 60)
        
        print(f"\n⏱️  总处理时间: {processing_time:.2f} 秒")
        
        return final_state.get('final_answer', '')
        
    def save_session_history(self):
        """保存会话历史到文件"""
        if not self.session_history:
            print("⚠️  没有会话历史可保存")
            return
            
        filename = f"session_history_{self.current_session_id}.json"
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(self.session_history, f, ensure_ascii=False, indent=2)
            print(f"💾 会话历史已保存到: {filename}")
        except Exception as e:
            print(f"❌ 保存会话历史失败: {e}")

    def show_help(self):
        """显示帮助信息"""
        help_text = """
🤖 多代理智能研究系统 - 交互模式帮助

基本命令:
  help          - 显示此帮助信息
  history       - 显示当前会话历史
  save          - 保存会话历史到文件
  clear         - 清空屏幕
  quit/exit/q   - 退出程序

多会话管理:
  sessions      - 列出所有会话
  new           - 创建新会话
  switch <id>   - 切换到指定会话
  rename <id> <name> - 重命名会话
  delete <id>   - 删除会话

功能:
  • 输入任何问题，系统将使用多个AI代理协作研究
  • 支持多会话管理，可以创建不同的对话上下文
  • 自动记录所有问题和答案
  • 可以保存完整的会话历史

示例问题:
  • "2024年巴黎奥运会有什么新的比赛项目？"
  • "解释量子计算的基本原理"
  • "写一篇关于人工智能伦理的文章"
        """
        print(help_text)
        
    def show_sessions(self):
        """显示所有会话"""
        if not self.sessions:
            print("📝 当前没有会话")
            return
            
        print("\n📝 所有会话")
        print("=" * 60)
        for session_id, session_data in self.sessions.items():
            current_marker = " 👈" if session_id == self.current_session_id else ""
            print(f"{session_data['name']} ({session_id}){current_marker}")
            print(f"   📅 创建时间: {session_data['created_at'][:19]}")
            print(f"   💬 消息数量: {len(session_data['history'])}")
            print()

    def show_history(self):
        """显示当前会话历史"""
        if not self.session_history:
            print("📝 当前没有会话历史")
            return
            
        print("\n📝 会话历史")
        print("=" * 60)
        for i, entry in enumerate(self.session_history, 1):
            print(f"{i}. [{entry['timestamp'][11:19]}] Q: {entry['question'][:50]}...")
            print(f"   ⏱️  {entry['processing_time']}秒 | ✅ 完成")
            print()

    def clear_screen(self):
        """清空屏幕"""
        os.system('cls' if os.name == 'nt' else 'clear')
        print("🤖 多代理智能研究系统 - 交互模式")
        print("=" * 60)

    def main(self):
        """主函数：交互式多代理工作流"""
        # 验证配置
        research_settings.model_validate(research_settings.model_dump())
        
        self.clear_screen()
        print("🤖 多代理智能研究系统 - 交互模式")
        print("=" * 60)
        print("输入您的问题，系统将使用多个AI代理协作研究")
        print("输入 'help' 查看可用命令")
        print("=" * 60)

        import asyncio
        asyncio.run(self._async_main())

    async def _async_main(self):
        """异步主循环"""
        while True:
            try:
                # 获取用户输入
                user_input = input("\n💬 请输入您的问题或命令: ").strip()
                
                # 检查退出条件
                if user_input.lower() in ['quit', 'exit', 'q']:
                    print("👋 感谢使用，再见！")
                    if self.session_history:
                        save_choice = input("💾 是否保存会话历史？(y/n): ").strip().lower()
                        if save_choice in ['y', 'yes']:
                            self.save_session_history()
                    break
                    
                # 检查命令
                if user_input.lower() == 'help':
                    self.show_help()
                    continue
                    
                if user_input.lower() == 'history':
                    self.show_history()
                    continue
                    
                if user_input.lower() == 'save':
                    self.save_session_history()
                    continue
                    
                if user_input.lower() == 'clear':
                    self.clear_screen()
                    continue
                    
                if user_input.lower() == 'sessions':
                    self.show_sessions()
                    continue
                    
                if user_input.lower() == 'new':
                    new_id = self.create_new_session()
                    print(f"✅ 已创建新会话: {self.sessions[new_id]['name']} ({new_id})")
                    continue
                    
                if user_input.lower().startswith('switch '):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        session_id = parts[1]
                        if self.switch_session(session_id):
                            print(f"✅ 已切换到会话: {self.sessions[session_id]['name']}")
                        else:
                            print("❌ 会话不存在")
                    else:
                        print("❌ 用法: switch <session_id>")
                    continue
                    
                if user_input.lower().startswith('rename '):
                    parts = user_input.split()
                    if len(parts) >= 3:
                        session_id = parts[1]
                        new_name = ' '.join(parts[2:])
                        if self.rename_session(session_id, new_name):
                            print(f"✅ 会话已重命名为: {new_name}")
                        else:
                            print("❌ 会话不存在")
                    else:
                        print("❌ 用法: rename <session_id> <new_name>")
                    continue
                    
                if user_input.lower().startswith('delete '):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        session_id = parts[1]
                        if self.delete_session(session_id):
                            print("✅ 会话已删除")
                        else:
                            print("❌ 无法删除会话（可能是最后一个会话或会话不存在）")
                    else:
                        print("❌ 用法: delete <session_id>")
                    continue
                
                # 检查空输入
                if not user_input:
                    print("⚠️  请输入有效的问题或命令")
                    continue
                    
                # 运行工作流处理问题
                await self.run_workflow(user_input)
                
            except KeyboardInterrupt:
                print("\n👋 程序被用户中断，再见！")
                if self.session_history:
                    save_choice = input("💾 是否保存会话历史？(y/n): ").strip().lower()
                    if save_choice in ['y', 'yes']:
                        self.save_session_history()
                break
            except Exception as e:
                print(f"❌ 处理过程中出现错误: {e}")
                import traceback
                traceback.print_exc()
                print("请重试或输入新问题")

if __name__ == "__main__":
    agent_system = InteractiveMultiAgent()
    agent_system.main()